# coding: UTF-8
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn import metrics
import time
from utils import get_time_dif
from tqdm import tqdm


# 权重初始化，默认xavier
def init_network(model, method='xavier', exclude='embedding', seed=123):
    for name, w in model.named_parameters():
        if exclude not in name:
            if 'weight' in name:
                if method == 'xavier':
                    nn.init.xavier_normal_(w)
                elif method == 'kaiming':
                    nn.init.kaiming_normal_(w)
                else:
                    nn.init.normal_(w)
            elif 'bias' in name:
                nn.init.constant_(w, 0)
            else:
                pass


def cross_entropy_to_mse(y_pred, y_true):
    y_pred_ = y_pred.view(-1)
    return 0.5 * torch.sum((y_true - y_pred_) ** 2) / len(y_true)


def train_1(config, model, train_iter, dev_iter, test_iter):
    start_time = time.time()
    model.train()
    optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)

    # 学习率指数衰减，每次epoch：学习率 = gamma * 学习率
    # scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9)
    total_batch = 0  # 记录进行到多少batch
    dev_best_loss = float('inf')
    last_improve = 0  # 记录上次验证集loss下降的batch数
    flag = False  # 记录是否很久没有效果提升
    for epoch in range(config.num_epochs):
        print('Epoch [{}/{}]'.format(epoch + 1, config.num_epochs))
        # scheduler.step() # 学习率衰减
        for i, (trains, labels) in enumerate(train_iter):
            outputs = model(trains)
            model.zero_grad()
            loss = F.cross_entropy(outputs, labels)
            loss.backward()
            optimizer.step()
            if total_batch % 100 == 0:
                # 每多少轮输出在训练集和验证集上的效果
                true = labels.data.cpu()
                predic = torch.max(outputs.data, 1)[1].cpu()
                train_acc = metrics.accuracy_score(true, predic)
                dev_acc, dev_loss = evaluate_1(config, model, dev_iter)
                if dev_loss < dev_best_loss:
                    dev_best_loss = dev_loss
                    torch.save(model.state_dict(), config.save_path)
                    improve = '*'
                    last_improve = total_batch
                else:
                    improve = ''
                time_dif = get_time_dif(start_time)
                msg = 'Iter: {0:>6},  Train Loss: {1:>5.2},  Train Acc: {2:>6.2%},  Val Loss: {3:>5.2},  Val Acc: {4:>6.2%},  Time: {5} {6}'
                print(msg.format(total_batch, loss.item(), train_acc, dev_loss, dev_acc, time_dif, improve))
                model.train()
            total_batch += 1
            if total_batch - last_improve > config.require_improvement:
                # 验证集loss超过1000batch没下降，结束训练
                print("No optimization for a long time, auto-stopping...")
                flag = True
                break
        if flag:
            break
    test(config, model, test_iter)


def train(config, model, train_iter, dev_iter, test_iter):
    start_time = time.time()
    model.train()
    optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)
    # 学习率指数衰减，每次epoch：学习率 = gamma * 学习率
    # scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9)
    total_batch = 0  # 记录进行到多少batch
    dev_best_loss = float('inf')
    last_improve = 0  # 记录上次验证集loss下降的batch数
    flag = False  # 记录是否很久没有效果提升
    for epoch in range(config.num_epochs):
        print('Epoch [{}/{}]'.format(epoch + 1, config.num_epochs))
        # scheduler.step() # 学习率衰减
        for i, (trains, labels) in tqdm(enumerate(train_iter)):
            outputs = model(trains)
            model.zero_grad()
            # labels = labels.to(outputs.dtype)
            loss = cross_entropy_to_mse(outputs, labels)
            loss.backward()
            optimizer.step()
            if total_batch % 10 == 0:
                dev_acc, dev_loss = evaluate(config, model, dev_iter)
                if dev_loss < dev_best_loss:
                    dev_best_loss = dev_loss
                    torch.save(model.state_dict(), config.save_path)
                    improve = '*'
                    last_improve = total_batch
                else:
                    improve = ''
                time_dif = get_time_dif(start_time)
                # msg = 'Iter: {0:>6},  Train Loss: {1:>5.2},  Train Acc: {2:>6.2%},  Val Loss: {3:>5.2},  Val Acc: {4:>6.2%},  Time: {5} {6}'
                msg = 'Iter: {0:>6},  Train Loss: {1:>5.2},  Val Loss: {2:>5.2},  Val Acc: {3:>6.2%},  Time: {4} {5}'
                # print(msg.format(total_batch, loss.item(), train_acc, dev_loss, dev_acc, time_dif, improve))
                print(msg.format(total_batch, loss.item(), dev_loss, dev_acc, time_dif, improve))
                model.train()
            total_batch += 1
            if total_batch - last_improve > config.require_improvement:
                # 验证集loss超过1000batch没下降，结束训练
                print("No optimization for a long time, auto-stopping...")
                flag = True
                break
        if flag:
            break
    test(config, model, test_iter)


def test_test(config, model, test_iter):
    # test
    model.load_state_dict(torch.load(config.save_path))
    print("Loading model...")
    start_time = time.time()
    model.eval()
    predict = evaluate_test(config, model, test_iter)
    zheng = 0
    total = 0
    for i in range(0, len(predict)):
        if predict[i] == 1 or predict[i] == 2:
            zheng = zheng + 1
        total = total + 1
    print("有知识输出占比", zheng * 1.0 / total)
    time_dif = get_time_dif(start_time)
    print("Time usage:", time_dif)


def test(config, model, test_iter, test_text):
    # test
    model.load_state_dict(torch.load(config.save_path))
    print("Loading model...")
    start_time = time.time()
    model.eval()
    test_acc, test_loss, test_report, test_confusion, predict_all = evaluate(config, model, test_iter, test_text, test=True)
    msg = 'Test Loss: {0:>5.2},  Test Acc: {1:>6.2%}'
    print(msg.format(test_loss, test_acc))
    print("Precision, Recall and F1-Score...")
    print(test_report)
    print("Confusion Matrix...")
    print(test_confusion)
    know_acc = (test_confusion[1][1] + test_confusion[1][2] + test_confusion[2][1] + test_confusion[2][2]) * 1.0 / (
            test_confusion[0][1] + test_confusion[0][2] + test_confusion[1][1] + test_confusion[1][2] +
            test_confusion[2][1] + test_confusion[2][2])
    print("有知识的准确", know_acc)
    zheng_recall = (test_confusion[1][1] + test_confusion[1][2] + test_confusion[2][1] + test_confusion[2][2]) * 1.0 / (
            test_confusion[1][0] + test_confusion[1][1] + test_confusion[1][2] + test_confusion[2][0] +
            test_confusion[2][1] + test_confusion[2][2])
    print("正样本召回率", zheng_recall)
    louzhao = (test_confusion[1][0] + test_confusion[2][0]) * 1.0 / (
            test_confusion[1][0] + test_confusion[1][1] + test_confusion[1][2] + test_confusion[2][0] +
            test_confusion[2][1] + test_confusion[2][2])
    print("正样本漏招率", louzhao)
    elie = test_confusion[2][0] * 1.0 / (test_confusion[2][0] + test_confusion[2][1] + test_confusion[2][2])
    print("恶劣漏招率", elie)
    knowledge = (test_confusion[1][1] + test_confusion[1][2] + test_confusion[2][1] + test_confusion[2][2] +
                 test_confusion[0][1] + test_confusion[0][2]) * 1.0 / 1000
    print("有知识的输出率", knowledge)
    time_dif = get_time_dif(start_time)
    print("Time usage:", time_dif)


def evaluate_1(config, model, data_iter, test=False):
    model.eval()
    loss_total = 0
    predict_all = np.array([], dtype=int)
    labels_all = np.array([], dtype=int)
    with torch.no_grad():
        for texts, labels in data_iter:
            outputs = model(texts)
            loss = F.cross_entropy(outputs, labels)
            loss_total += loss
            labels = labels.data.cpu().numpy()
            predic = torch.max(outputs.data, 1)[1].cpu().numpy()
            labels_all = np.append(labels_all, labels)
            predict_all = np.append(predict_all, predic)
    acc = metrics.accuracy_score(labels_all, predict_all)
    if test:
        report = metrics.classification_report(labels_all, predict_all, target_names=config.class_list, digits=4)
        confusion = metrics.confusion_matrix(labels_all, predict_all)
        return acc, loss_total / len(data_iter), report, confusion
    return acc, loss_total / len(data_iter)


def evaluate_test(config, model, data_iter):
    model.eval()
    predict_all = np.array([], dtype=int)
    with torch.no_grad():
        for texts in data_iter:
            outputs = model(texts)
            predic = []
            for k in range(0, len(outputs.data)):
                if outputs.data[k][0] < 0.5:
                    predic.append(0)
                elif 0.5 <= outputs.data[k][0] <= 1.5:
                    predic.append(1)
                else:
                    predic.append(2)
            predic = torch.tensor(np.array(predic)).cpu()
            predict = predic.numpy()
            predict_all = np.append(predict_all, predict)

    return predict_all


def evaluate(config, model, data_iter, test_text, test=False):
    model.eval()
    loss_total = 0
    predict_all = np.array([], dtype=int)
    labels_all = np.array([], dtype=int)
    # text_1 = []
    text_2 = []
    m = 0
    with torch.no_grad():
        for texts, labels in data_iter:
            outputs = model(texts)
            loss = cross_entropy_to_mse(outputs, labels)
            # print("outputs", outputs)
            predic = []
            for k in range(0, len(outputs.data)):
                if outputs.data[k][0] < 0.8:
                    if int(labels[k]) == 2:
                        print(test_text[m])
                        print(labels[k])
                        text_2.append(test_text[m])
                    # if int(labels[k]) == 1:
                    #     text_1.append(test_text[k])
                    predic.append(0)
                elif 0.8 <= outputs.data[k][0] <= 1.6:
                    predic.append(1)
                else:
                    predic.append(2)
                m = m + 1
            predic = torch.tensor(np.array(predic)).cpu()
            labels = labels.data.cpu()
            # loss = F.cross_entropy(predic, labels)
            loss_total += loss
            # predic = torch.max(outputs.data, 1)[1].cpu().numpy()

            predict = predic.numpy()
            labels_all = np.append(labels_all, labels)
            predict_all = np.append(predict_all, predict)
    with open('case_8.txt', 'w', encoding='utf-8') as f:
        for text in text_2:
            f.write(text + '\n')
        f.close()
    # with open('case_1.txt', 'w', encoding='utf-8') as f:
    #     for text in text_1:
    #         f.write(text + '\n')
    #     f.close()
    # print("text_1", text_1)
    # print("text_2", text_2)
    # print(len(text_1))
    # print(len(text_2))
    acc = metrics.accuracy_score(labels_all, predict_all)
    if test:
        report = metrics.classification_report(labels_all, predict_all, target_names=config.class_list, digits=4)
        confusion = metrics.confusion_matrix(labels_all, predict_all)
        return acc, loss_total / len(data_iter), report, confusion, predict_all
    return acc, loss_total / len(data_iter)
